import py_trees
import time

# 定义几个简单的行为节点
class CountTask(py_trees.behaviour.Behaviour):
    """计数任务，模拟耗时操作"""
    def __init__(self, name: str, count: int):
        super().__init__(name)
        self.count = count
        self.current = 0

    def update(self) -> py_trees.common.Status:
        self.current += 1
        print(f"[{self.name}] 计数: {self.current}/{self.count}")
        
        if self.current >= self.count:
            print(f"[{self.name}] 完成!")
            return py_trees.common.Status.SUCCESS
        else:
            return py_trees.common.Status.RUNNING

    def terminate(self, new_status: py_trees.common.Status):
        if new_status == py_trees.common.Status.INVALID:
            print(f"[{self.name}] 被中断")


def create_parallel_tree():
    """创建包含并行节点的行为树"""
    # 创建三个不同计数任务的子节点
    task1 = CountTask("任务1", 3)  # 需要3步完成
    task2 = CountTask("任务2", 5)  # 需要5步完成
    task3 = CountTask("任务3", 2)  # 需要2步完成

    # 创建并行节点：
    # 策略1: SuccessOnAll() - 所有子节点成功才返回成功
    # 策略2: SuccessOnOne() - 任一子节点成功就返回成功
    parallel = py_trees.composites.Parallel(
        name="并行执行",
        policy=py_trees.common.ParallelPolicy.SuccessOnAll()
    )

    # 添加子节点到并行节点
    parallel.add_children([task1, task2, task3])

    return parallel


if __name__ == "__main__":
    # 创建行为树
    behavior_tree = create_parallel_tree()

    # 打印树结构
    print("行为树结构:")
    print(py_trees.display.ascii_tree(behavior_tree))

    # 执行行为树
    print("\n开始执行并行任务...")
    start_time = time.time()
    for _ in range(10):  # 最多执行10步
        behavior_tree.tick_once()  # 执行一次树迭代
        if behavior_tree.status == py_trees.common.Status.SUCCESS:
            print("\n所有并行任务完成!")
            break
        print("----------------->")
        time.sleep(1)  # 等待一下，方便观察输出
    
    print(f"\n执行时间: {time.time() - start_time:.2f}秒")